#include "yolov8_tensorrt.h"
#include <opencv2/opencv.hpp>
#include <vector>
#include <string>
#include <iostream>

// 只初始化一次
YOLOv8TensorRT& get_classifier() {
    static YOLOv8TensorRT classifier("/home/nano/code/yolo/test_number/best_class.engine");
    return classifier;
}

// 分类结果函数 - 返回 0,1,2,3
// int classify_image(const cv::Mat& img) {
//     cv::Mat image = img.clone();
//     YOLOv8TensorRT& classifier = get_classifier();

//     // 进行推理
//     std::vector<float> output = classifier.classify(image);
    
//     // 找到最大概率的类别
//     int max_idx = 0;
//     float max_prob = output[0];print();
    
//     for (int i = 1; i < output.size(); i++) {
//         if (output[i] > max_prob) {
//             max_prob = output[i];
//             max_idx = i;
//         }
//     }
    
//     // 如果最大概率低于阈值，返回0（背景）
//     if (max_prob < 0.65f) { // 可以根据需要调整阈值
//         return 0;
//     }
    
//     // 返回类别索引（1,2,3）
//     return max_idx;
// }

// 如果需要保持原有函数名但修改功能

int funtion(cv::Mat& img) {

    // cv::Mat image = img.clone();
    YOLOv8TensorRT& classifier = get_classifier();
    // 进行推理
    std::vector<float> output = classifier.classify(img);
    
    // 找到最大概率的类别
    int max_idx = 0;
    float max_prob = output[0];
    
    for (int i = 1; i < output.size(); i++) {
        if (output[i] > max_prob) {
            max_prob = output[i];
            max_idx = i;
        }
    }
    int result = 0;
    // 如果最大概率低于阈值，返回0（背景）
    if (max_prob >= ODDS) { // 可以根据需要调整阈值
        result = max_idx;
    }
    // int result = classify_image(img);
    
    // 创建一个显示结果的图像（可选）
    // cv::Mat result_image = img.clone();
    cvtColor(img,img,cv::COLOR_GRAY2BGR);
    int fontFace = cv::FONT_HERSHEY_SIMPLEX;
    double fontScale = 0.7;
    int thickness = 2;
    int margin = 10;
    int lineHeight = 30;
    // 创建标签文本
    std::string label = "Class: " + std::to_string(result);
    // 在图像上绘制标签
                // 想要绘制的标签，文本内容
    cv::putText(img, label,
        //      标签位置 右边上角   
               cv::Point(img.cols - 150, 30),
               cv::FONT_HERSHEY_SIMPLEX, 1.0,
               cv::Scalar(0, 200, 200), 2);



    //计算文本位置
    // int startY = result_image.rows - margin - (lineHeight * 3); // 3行文本
    // for(int i = 0; i < 3; i++)
    // {
    //     std::string probText = "Class " + std::to_string(i+1) + ": "+
    //         std::to_string(output[i+1]).substr(0, 4); // 保留4位小数
        
        // cv::Point textOrg(margin, startY + i * lineHeight);
    for(int i = 0; i < output.size(); i++)
    {
        std::string probText = "Class " + std::to_string(i) + ": "+
            std::to_string(output[i]).substr(0, 4); // 保留4位小数

        cv::putText(img,
                    probText,
                    // textOrg,
                    cv::Point(img.cols - 150, 60 + i * lineHeight),
                    fontFace,
                    fontScale,
                    cv::Scalar(255, 0, 0),
                    thickness);
    }
    // imshow("模型后的结果：", result_image);
    // cv::waitKey(1); // 确保图像窗口刷新
    return result;
}